import subprocess
import argparse
import time
import sys
import os
from datetime import datetime

# 显式地设置标准输出和标准错误输出的编码为UTF-8
#sys.stdout = open(sys.stdout.fileno(), mode='w', encoding='utf-8', buffering=1)
#sys.stderr = open(sys.stderr.fileno(), mode='w', encoding='utf-8', buffering=1)

def check_device_connected():
    # 检查设备是否存在
    check_device_command = "hdc list targets"
    result = subprocess.run(check_device_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    return "[Empty]" not in result.stdout

def delete_file(filename):
    current_dir = os.getcwd()
    file_path = os.path.join(current_dir, filename)    
    if os.path.exists(file_path):
        os.remove(file_path)

def delete_snapshot_on_device():
    subprocess.run(['hdc', 'shell', 'rm -f /data/local/tmp/screenshot.jpeg'])        

def write_reboot_record(reboot_no,reboot_time, success):
    with open('reboot_record.txt', 'a') as f:
        f.write(f"{reboot_no} - {reboot_time} - {'Success' if success else 'Failure'}\n")

def restart_device(restart_count):
    delete_file("reboot_record.txt")    
    for i in range(1, restart_count + 1):
        if not check_device_connected():
            print("******没有检测到设备******", flush=True)
            break        
        print(f"******计划重启{restart_count}次，第 {i} 次重启******",flush=True)
        # hdc命令用于重启设备
        command = "hdc shell reboot"
        reboot_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        reboot_no = i
        try:
            # 执行hdc命令
            subprocess.run(command, shell=True, check=True)
            print("设备正在重启...",flush=True)
        except subprocess.CalledProcessError as e:
            print(f"重启设备失败: {e}",flush=True)
            continue

        # 等待设备重启并进入idle界面
        print("等待设备重启并进入idle界面...",flush=True)
        time.sleep(40)  # 假设设备重启并进入idle界面需要大约40秒

        # 在尝试截图前再次检查设备是否存在
        if not check_device_connected():
            print("******没有检测到设备******", flush=True)
            write_reboot_record(reboot_no,reboot_time, False)
            break
        # 尝试截图
        screenshot_command = "hdc shell snapshot_display -f /data/local/tmp/screenshot.jpeg"
        pull_screenshot_command = "hdc file recv /data/local/tmp/screenshot.jpeg"
       
        try:
            # 执行截图命令
            subprocess.run(screenshot_command, shell=True, check=True)
            # 将截图从设备拉取到本地
            subprocess.run(pull_screenshot_command, shell=True, check=True)
            print("设备重启成功，并且截图已保存到本地",flush=True)
            write_reboot_record(reboot_no,reboot_time, True)            
            delete_file("screenshot.jpeg")
            delete_snapshot_on_device()
        except subprocess.CalledProcessError as e:
            print(f"设备重启后截图失败: {e}",flush=True)
            write_reboot_record(reboot_no,reboot_time, False)            

if __name__ == "__main__":
    # 创建 ArgumentParser 对象
    parser = argparse.ArgumentParser(description='重启鸿蒙设备的脚本。')

    # 添加参数
    parser.add_argument('count', type=int, help='重启设备的次数')

    # 解析参数
    args = parser.parse_args()

    # 调用重启函数
    restart_device(args.count)
